package cn.edu.hnu.matching.websocket;

import cn.edu.hnu.matching.toolcalling.AiChatTool;
import cn.hutool.json.JSONObject;
import cn.hutool.json.JSONUtil;
import lombok.RequiredArgsConstructor;
import org.springframework.ai.chat.client.ChatClient;
import org.springframework.ai.chat.prompt.ChatOptions;
import org.springframework.ai.chat.prompt.Prompt;
import org.springframework.stereotype.Component;
import org.springframework.web.socket.TextMessage;
import org.springframework.web.socket.WebSocketSession;
import org.springframework.web.socket.handler.TextWebSocketHandler;

import java.io.IOException;

@Component
@RequiredArgsConstructor
public class AiWebSocketHandler extends TextWebSocketHandler {

    private final ChatClient chatClient;

    private final AiChatTool aiChatTool;

    @Override
    public void handleTextMessage(WebSocketSession session, TextMessage message) throws IOException {
        String json = message.getPayload();
        JSONObject data = JSONUtil.parseObj(json);

        String originalPrompt = data.getStr("prompt");
        String token = data.getStr("token");

        String prompt = originalPrompt + "(当前用户登录 token: %s)".formatted(token);

        chatClient.prompt(new Prompt(
                    prompt,
                    ChatOptions
                        .builder()
                        .model("deepseek-v3")
                        .build()
                ))
                .tools(aiChatTool)
                .stream()
                .content()
                .doOnNext(part -> {
                    try {
                        session.sendMessage(new TextMessage(part));
                    } catch (IOException e) {
                        e.printStackTrace();
                    }
                })
                .doOnComplete(() -> {
                    try {
                        session.sendMessage(new TextMessage("[complete]"));
                        session.close();
                    } catch (IOException e) {
                        e.printStackTrace();
                    }
                })
                .subscribe();
    }
}
